# -*- coding: utf-8 -*-

from functools import wraps

from marshmallow import Schema


def marshal_with(mallow, envelope=None):
    assert issubclass(mallow, Schema)

    def decorator(f):
        @wraps(f)
        def wrapper(*args, **kwargs):
            resp = mallow().dump(f(*args, **kwargs))
            if envelope:
                resp = {envelope: resp}
            return resp
        return wrapper
    return decorator
